package com.kingbook.core.executor;

import com.kingbook.common.utils.StringUtils;
import com.kingbook.core.utils.JDBCHelper;
import com.google.common.collect.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.*;
import org.springframework.jdbc.core.namedparam.*;
import org.springframework.jdbc.support.GeneratedKeyHolder;
import org.springframework.jdbc.support.KeyHolder;
import org.springframework.lang.Nullable;

import java.sql.*;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * @Author: Bucky
 * @Date: 2019/1/2
 * @Version 1.0
 */
public class ParsedNamedParameterJdbcTemplate extends NamedParameterJdbcTemplate {

    private static Logger logger = LoggerFactory.getLogger(ParsedNamedParameterJdbcTemplate.class);

    ParsedNamedParameterJdbcTemplate(JdbcOperations classicJdbcTemplate) {
        super(classicJdbcTemplate);
    }

    @Override
    public <T> T query(String sql, Map<String, ?> paramMap, ResultSetExtractor<T> rse) {
        if (paramMap == null || paramMap.isEmpty())
            return getJdbcTemplate().query(sql, rse);
        else
            return super.query(sql, paramMap, rse);
    }

    public <T> List<T> queryList(String sql, Map<String, ?> paramMap, BeanListResultSetExtractor<T> rse) {
        if (paramMap == null || paramMap.isEmpty())
            return getJdbcTemplate().query(sql, rse);
        else
            return super.query(sql, paramMap, rse);
    }

    public List<String> extractTableNames(String sql, Map<String, ?> paramMap) {
        sql = "EXPLAIN " + sql;
        if (paramMap == null || paramMap.isEmpty())
            return getJdbcTemplate().query(sql,new TablesResultSetExtractor());
        else
            return super.query(sql, paramMap, new TablesResultSetExtractor());
    }

    @Override
    public int update(String sql, Map<String, ?> paramMap) {
        if (paramMap == null || paramMap.isEmpty())
            return getJdbcTemplate().update(sql);
        else
            return super.update(sql, paramMap);
    }

    public int insertReturnKey(final String sql, Map<String, ?> paramMap) {
        KeyHolder keyHolder = new GeneratedKeyHolder();
        int rows = 0;
        if (paramMap == null || paramMap.isEmpty())
            rows = getJdbcTemplate().update((con) -> {
                return con.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS);
            }, keyHolder);
        else
            rows = new NamedParameterJdbcTemplate(getJdbcTemplate()).update(sql, new MapSqlParameterSource(paramMap), keyHolder);

        logger.debug("key:" + keyHolder.getKeyList());
        if (rows == 1)
            return keyHolder.getKey().intValue();
        else if (rows > 1) {
            List<Map<String, Object>> list = keyHolder.getKeyList();
            int[] ids = new int[list.size()];
            for (int j = 0; j < ids.length; j++)
                ids[j] = Integer.parseInt(String.valueOf(list.get(j).get("GENERATED_KEY")));
            return ids[0];
        }
        return -1;
    }

    @Override
    @Nullable
    public int[] batchUpdate(@Nullable String sql, Map<String, ?>[] batchValues) {
        if (batchValues == null || batchValues.length == 0)
            return getJdbcTemplate().batchUpdate(sql);
        else {
            SqlParameterSource[] batchArgs = new SqlParameterSource[batchValues.length];
            int i = 0;
            for (Map<String, ?> values : batchValues) {
                batchArgs[i] = new MapSqlParameterSource(values);
                i++;
            }
            try {
                ParsedSql parsedSql = getParsedSql(sql);
                Connection con = getJdbcTemplate().getDataSource().getConnection();
                return new BatchSetterUtils(con).executeBatchUpdateWithNamedParameters(parsedSql, batchArgs, getJdbcOperations());
            } catch (SQLException e) {
                logger.error(e.getLocalizedMessage());
                return new int[]{-1};
            }
        }
    }

    int[] batchUpdate(String sql, List<Map<String, ?>> batchValues) {
        if (batchValues == null || batchValues.isEmpty())
            return getJdbcTemplate().batchUpdate(sql);
        else {
            SqlParameterSource[] batchArgs = new SqlParameterSource[batchValues.size()];
            int i = 0;
            for (Map<String, ?> values : batchValues) {
                batchArgs[i] = new MapSqlParameterSource(values);
                i++;
            }
            try {
                ParsedSql parsedSql = getParsedSql(sql);
                Connection con = getJdbcTemplate().getDataSource().getConnection();
                return new BatchSetterUtils(con).executeBatchUpdateWithNamedParameters(parsedSql, batchArgs, getJdbcOperations());
            } catch (SQLException e) {
                logger.error(e.getLocalizedMessage());
                return new int[]{-1};
            }
        }
    }


    class BatchSetterUtils extends BatchUpdateUtils {
        Connection con;
        BatchSetterUtils(Connection connection){
            con = connection;
        }
        int[] executeBatchUpdateWithNamedParameters(final ParsedSql parsedSql, final SqlParameterSource[] batchArgs, JdbcOperations jdbcOperations) {
            if (batchArgs.length <= 0) {
                return new int[]{0};
            }
            String sqlToUse = NamedParameterUtils.substituteNamedParameters(parsedSql, batchArgs[0]);
            return new BacthOperation(con).batchUpdate(sqlToUse, new BatchPreparedStatementSetter() {
                @Override
                public void setValues(PreparedStatement ps, int i) throws SQLException {
                    Object[] values = NamedParameterUtils.buildValueArray(parsedSql, batchArgs[i], null);
                    int[] columnTypes = NamedParameterUtils.buildSqlTypeArray(parsedSql, batchArgs[i]);
                    setStatementParameters(values, ps, columnTypes);
                }

                @Override
                public int getBatchSize() {
                    return batchArgs.length;
                }
            });
        }
    }

    class BacthOperation extends JdbcTemplate {
        Connection con;
        BacthOperation(Connection connection){
            con = connection;
        }
        @Override
        @Nullable
        public int[] batchUpdate(final String sql, final BatchPreparedStatementSetter pss) throws DataAccessException {
            PreparedStatement ps = null;
            ResultSet rs = null;
            try {
                ps = con.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS);
                int batchSize = pss.getBatchSize();
                for (int i = 0; i < batchSize; i++) {
                    pss.setValues(ps, i);
                    ps.addBatch();
                }
                ps.executeBatch();
                rs = ps.getGeneratedKeys();
                List<Map<String, Object>> list = Lists.newArrayList();
                if (rs != null) {
                    RowMapperResultSetExtractor<Map<String, Object>> rse = new RowMapperResultSetExtractor<Map<String, Object>>(getColumnMapRowMapper(), 1);
                    list = rse.extractData(rs);
                }
                int[] ids = new int[list.size()];
                for (int j = 0; j < ids.length; j++)
                    ids[j] = Integer.parseInt(String.valueOf(list.get(j).get("GENERATED_KEY")));
                return ids;
            } catch (SQLException ex) {
                throw getExceptionTranslator().translate("PreparedStatementCallback", sql, ex);
            } catch (Exception e) {
                logger.error(e);
                return new int[]{-1};
            } finally {
                JDBCHelper.cleanup(con, ps, rs);
            }
        }
    }


    class TablesResultSetExtractor implements ResultSetExtractor<List<String>> {
        @Override
        public List<String> extractData(ResultSet rs) throws SQLException {
            List<String> tables = Lists.newArrayList();
            while (rs.next()) {
                tables.add(rs.getString("table"));
            }
            return tables;
        }

    }

}
